import math
import os
import sys
import traceback
from functools import wraps
from typing import *

import numpy as np
import pandas as pd
from matplotlib import pyplot as plt
from sklearn.metrics import f1_score

from utils.fscore_utils import *


# consider delay threshold and missing segments
def get_range_proba(predict, label, delay=7):
    splits = np.where(label[1:] != label[:-1])[0] + 1
    is_anomaly = label[0] == 1
    new_predict = np.array(predict)
    pos = 0

    for sp in splits:
        if is_anomaly:
            if 1 in predict[pos:min(pos + delay + 1, sp)]:
                new_predict[pos: sp] = 1
            else:
                new_predict[pos: sp] = 0
        is_anomaly = not is_anomaly
        pos = sp
    sp = len(label)

    if is_anomaly:  # anomaly in the end
        if 1 in predict[pos: min(pos + delay + 1, sp)]:
            new_predict[pos: sp] = 1
        else:
            new_predict[pos: sp] = 0

    return new_predict


# set missing = 0
def reconstruct_label(timestamp, label):
    timestamp = np.asarray(timestamp, np.int64)
    index = np.argsort(timestamp)

    timestamp_sorted = np.asarray(timestamp[index])
    interval = np.min(np.diff(timestamp_sorted))

    label = np.asarray(label, np.int64)
    label = np.asarray(label[index])

    idx = (timestamp_sorted - timestamp_sorted[0]) // interval

    new_label = np.zeros(shape=((timestamp_sorted[-1] - timestamp_sorted[0]) // interval + 1,), dtype=np.int)
    new_label[idx] = label

    return new_label


def fscore_for_threshold(threshold: float, nll_list: np.ndarray, label_list: np.ndarray):
    pred_list = (nll_list > threshold).astype(np.int64)
    y_pred = get_range_proba(pred_list, label_list, 7)
    fscore = f1_score(label_list, y_pred)

    return fscore


def get_best_fscore(ts_list: np.ndarray,
                    nll_list: np.ndarray,
                    label_list: np.ndarray):
    nll_list = reconstruct_label(ts_list, nll_list)
    label_list = reconstruct_label(ts_list, label_list)

    l, r = nll_list.min(), nll_list.max()

    for _ in range(20):
        m = (l + r) / 2
        mm = (m + r) / 2

        val1 = fscore_for_threshold(m, nll_list, label_list)
        val2 = fscore_for_threshold(mm, nll_list, label_list)

        if val1 < val2:
            l = m
        else:
            r = mm

    threshold = (l + r) / 2

    result = {
        'best_threshold': threshold,
        'best_fscore': fscore_for_threshold(threshold, nll_list, label_list),
        'nll_normal': float(np.mean(nll_list[label_list == 0])),
        'nll_anomaly': float(np.mean(nll_list[label_list == 1]))
    }

    return result


def analyze_anomaly_nll(nll_list: np.ndarray,
                        label_list: np.ndarray,
                        up_sample_normal: int = 1
                        ) -> Dict[str, float]:

    def log_error(method, default_value=None):
        @wraps(method)
        def wrapper(*args, **kwargs):
            try:
                return method(*args, **kwargs)
            except Exception:
                print(''.join(traceback.format_exception(*sys.exc_info())), file=sys.stderr)
                return default_value
        return wrapper

    def call_plot(fn_, *args, output_file, **kwargs):
        if output_file == ':show:':
            fig = fn_(*args, **kwargs)
            plt.show()
            plt.close()
        else:
            fn_(*args, output_file=output_file, **kwargs)

    # up sample normal nll & label if required
    if up_sample_normal and up_sample_normal > 1:
        normal_nll = nll_list[label_list == 0]
        normal_label = label_list[label_list == 0]
        nll_list = np.concatenate(
            [normal_nll] * (up_sample_normal - 1) + [nll_list],
            axis=0
        )
        label_list = np.concatenate(
            [normal_label] * (up_sample_normal - 1) + [label_list],
            axis=0
        )

    # prepare for analyze
    result_dict = {}
    is_anomaly_list = label_list != 0

    # separated nlls for different labels
    result_dict['nll_normal'] = float(np.mean(nll_list[label_list == 0]))
    result_dict['nll_anomaly'] = float(np.mean(nll_list[label_list == 1]))

    # auc score
    result_dict['auc'] = float(auc_score(nll_list, is_anomaly_list))

    # best f-score
    F = log_error(best_fscore, default_value=(math.nan, math.nan))

    def best_fscore_for_label(label):
        not_label = 2 if label == 1 else 1
        mask = label_list != not_label
        return F(nll_list[mask], label_list[mask] != 0)

    best_fscore_total, best_threshold, best_pr_total, best_rc_total = F(nll_list, is_anomaly_list)
    result_dict.update({
        'best_fscore': float(best_fscore_total),
        'best_threshold': float(best_threshold)
    })

    return result_dict
